#!/usr/bin/env python3

import argparse
from glob import glob
import os
import pendulum
from jinja2 import Template
import re
from shutil import copy
import sys
script_root = os.path.dirname(os.path.realpath(__file__))
sys.path.append(f'{script_root}/../utils')
from utils import cli, has_key, check_files, edit_file, run, parse_config, submit_job

def run_wps_geogrid(work_root, wps_root, config, args):
	wps_work_dir = os.path.abspath(work_root) + '/wps'
	if not os.path.isdir(wps_work_dir): os.mkdir(wps_work_dir)
	old_dir = os.getcwd()
	os.chdir(wps_work_dir)

	cli.notice(f'Run geogrid.exe at {wps_work_dir} ...')
	if not os.path.isfile('GEOGRID.TBL'):
		run(f'cp {wps_root}/geogrid/GEOGRID.TBL.ARW {wps_work_dir}/GEOGRID.TBL')
	if has_key(config, ('custom', 'GEOGRID.TBL', 'smooth_passes')):
		edit_file('GEOGRID.TBL', [
			['smooth_passes=.*', f'smooth_passes={config["custom"]["GEOGRID.TBL"]["smooth_passes"]}']
		])
	expected_files = ['geo_em.d{:02d}.nc'.format(i + 1) for i in range(config['domains']['max_dom'])]
	if not check_files(expected_files):
		run('rm -f geo_em.d*.nc')
		submit_job(f'{wps_root}/geogrid/src/geogrid.exe', min(args.np, args.np), config, args, logfile='geogrid.log.0000', wait=True)
		if not check_files(expected_files):
			cli.error(f'Failed! Check output {os.path.abspath(wps_work_dir)}/geogrid.log.0000')
		run(f'{script_root}/../plots/plot_domains.py -i geo_em.d*')
		cli.notice('Succeeded.')
	else:
		cli.notice('File geo_em.*.nc already exist.')
	run(f'ls -l {wps_work_dir}/geo_em.*.nc')
	os.chdir(old_dir)
